import json
from datetime import datetime
import argparse
import os


def json_to_markdown_tree(json_file):
  with open(json_file, 'r', encoding='utf-8') as f:
    data = json.load(f)

  nodes = data['nodes']
  links = data['links']

  # 根据y坐标排序节点
  sorted_nodes = sorted(enumerate(nodes), key=lambda x: x[1]['pos'][1])
  node_order = {node_id: i for i, (node_id, _) in enumerate(sorted_nodes)}

  children = {i: [] for i in range(len(nodes))}
  for link in links:
    children[link['src']].append(link['tgt'])

  root_nodes = set(range(len(nodes))) - set(link['tgt'] for link in links)

  def build_tree(node_id, prefix=''):
    node = nodes[node_id]
    tree = f"{prefix}├── {node['name']}\n"
    if node['text']:
      tree += f"{prefix}│   └── {node['text']}\n"

    sorted_children = sorted(children[node_id], key=lambda x: node_order[x])
    for i, child in enumerate(sorted_children):
      if i == len(sorted_children) - 1:
        tree += build_tree(child, prefix + "│   ")
      else:
        tree += build_tree(child, prefix + "│   ")

    return tree

  markdown = f"# 小说梳理\n\n"
  markdown += f"创建时间：{datetime.now().strftime('%Y-%m-%d %H:%M:%S')}\n"
  markdown += f"版本：{data['version']}\n\n"
  markdown += "```\n"  # 开始代码块，以保持格式

  sorted_roots = sorted(root_nodes, key=lambda x: node_order[x])
  for i, root in enumerate(sorted_roots):
    if i == len(sorted_roots) - 1:
      markdown += build_tree(root)
    else:
      markdown += build_tree(root)

  independent_nodes = set(range(len(nodes))) - root_nodes - set(
    child for children_list in children.values() for child in children_list)
  if independent_nodes:
    markdown += "其他元素\n"
    for node_id in sorted(independent_nodes, key=lambda x: node_order[x]):
      node = nodes[node_id]
      markdown += f"├── {node['name']}\n"
      if node['text']:
        markdown += f"│   └── {node['text']}\n"

  markdown += "```\n"  # 结束代码块

  return markdown


def main():
    parser = argparse.ArgumentParser(description='将JSON文件转换为Markdown树状图')
    parser.add_argument('input', help='输入JSON文件的路径')
    parser.add_argument('-o', '--output', help='输出Markdown文件的路径（可选）')
    args = parser.parse_args()

    input_file = args.input
    if not os.path.exists(input_file):
        print(f"错误：输入文件 '{input_file}' 不存在")
        return

    markdown_output = json_to_markdown_tree(input_file)

    if args.output:
        output_file = args.output
    else:
        output_file = os.path.splitext(input_file)[0] + '_树状图.md'

    with open(output_file, 'w', encoding='utf-8') as f:
        f.write(markdown_output)

    print(f"Markdown文件已生成: {output_file}")


if __name__ == '__main__':
    main()